/**
* @brief: 定义了一些用于存储 yolt 网络输出参数的类
* @author: mysticalwing
* @modify: name            time            brief
*          mysticalwing    2018.11.10      create
*          mysticalwing    2018.11.11      delete class YoltBoundingBoxParams YoltGridCell, and rewrite YoltLabel
*/

#ifndef YOLT_DEF_H
#define YOLT_DEF_H

#include <vector>
#include <cassert>
#include <cstring>
#include <fstream>
#include <algorithm>
#include "./tiny_dnn/tiny_dnn.h"

#define YOLT_DEF_TEST false

#include <iostream>

#if YOLT_DEF_TEST

using std::cout;
using std::cin;
using std::endl;

#endif

namespace nn_yolt
{
typedef std::string str_t;

typedef tiny_dnn::float_t float_t;

typedef std::vector<int> i_vec_t;

typedef tiny_dnn::vec_t f_vec_t;

typedef std::vector<str_t> s_vec_t;

typedef std::vector<bool> b_vec_t;

typedef std::vector<b_vec_t> connection_table_t;

/****************************************************************************************
 * @brief: 存放 yolt 标签的类
 */
class YoltLabel
{
#define PARAM_NUM_PER_BOX 5

  private:
    int grid_cell_ver_num; // 垂直方向划分的 grid_cll 的个数
    int grid_cell_hor_num; // 水平方向划分的 grid_cll 的个数
    int box_num_per_cell;  // 每个 grid_cell 预测的 bounding box 的个数
    int cell_width;
    int cell_height;
    int unit_size;
    int row_size;
    float_t *values;

  public:
    YoltLabel(const int image_width,
              const int image_height,
              const int _grid_cell_hor_num,
              const int _grid_cell_ver_num,
              const int _box_num_per_cell)
        : grid_cell_ver_num(_grid_cell_ver_num),
          grid_cell_hor_num(_grid_cell_hor_num),
          box_num_per_cell(_box_num_per_cell)
    {
        assert((image_width % _grid_cell_hor_num) == 0);
        assert((image_height % _grid_cell_ver_num) == 0);

        cell_width = image_width / _grid_cell_hor_num;
        cell_height = image_height / _grid_cell_ver_num;

        unit_size = _box_num_per_cell * PARAM_NUM_PER_BOX;
        row_size = unit_size * _grid_cell_hor_num;

        values = new float_t[_grid_cell_ver_num * row_size];
        std::memset(values, 0.0f, sizeof(float_t) * _grid_cell_ver_num * row_size);
    }

    YoltLabel(const YoltLabel &right)
        : grid_cell_ver_num(right.grid_cell_ver_num),
          grid_cell_hor_num(right.grid_cell_hor_num),
          box_num_per_cell(right.box_num_per_cell),
          cell_width(right.cell_width), cell_height(right.cell_height),
          unit_size(right.unit_size), row_size(right.row_size)
    {
        values = new float_t[right.grid_cell_ver_num * right.row_size];
        std::memcpy(values, right.values, right.grid_cell_ver_num * right.row_size * sizeof(float_t));
    }

    ~YoltLabel()
    {
        delete[] values;
    }

    YoltLabel &operator=(const YoltLabel &right)
    {
        grid_cell_ver_num = right.grid_cell_ver_num;
        grid_cell_hor_num = right.grid_cell_hor_num;
        box_num_per_cell = right.box_num_per_cell;
        cell_width = right.cell_width;
        cell_height = right.cell_height;
        unit_size = right.unit_size;
        row_size = right.row_size;

        std::memcpy(values, right.values, right.grid_cell_ver_num * right.row_size * sizeof(float_t));

        return *this;
    }

    void clear()
    {
        std::memset(values, 0.0f, sizeof(float_t) * grid_cell_ver_num * row_size);
    }

    f_vec_t get(const int cell_nx, const int cell_ny, const int box_num) const
    {
        const float_t *p = this->operator()(cell_nx, cell_ny, box_num);
        return f_vec_t(p, p + PARAM_NUM_PER_BOX);
    }

    void set(const int cell_nx, const int cell_ny, const int box_num, const f_vec_t &set_value)
    {
        assert(set_value.size() == PARAM_NUM_PER_BOX);

        float_t *p = this->operator()(cell_nx, cell_ny, box_num);
        for (int h = 0; h < PARAM_NUM_PER_BOX; h++)
            p[h] = set_value[h];
    }

    void set(const int cell_nx, const int cell_ny, const int box_num,
             const float_t x, const float_t y,
             const float_t w, const float_t h, const float_t iou)
    {
        float_t *p = this->operator()(cell_nx, cell_ny, box_num);
        p[0] = x;
        p[1] = y;
        p[2] = w;
        p[3] = h;
        p[4] = iou;
    }

    // 将数据保存到指定文件中
    void save(const str_t &file_name) const
    {
        std::ofstream f_out;
        f_out.open(file_name.c_str(), std::ios_base::trunc | std::ios_base::out);

        assert(f_out.is_open());

        for (int i = 1; i <= row_size * grid_cell_ver_num; i++)
        {
            f_out << values[i - 1] << " ";
            if ((i % PARAM_NUM_PER_BOX) == 0)
            {
                f_out << std::endl;
            }
        }

        f_out.close();
    }

    // 从指定文件中读取数据
    void load(const str_t file_name)
    {
        clear();

        std::ifstream f_in;
        f_in.open(file_name.c_str());

        assert(f_in.is_open());

        for (int i = 0; i < row_size * grid_cell_ver_num; i++)
        {
            f_in >> values[i];
        }

        f_in.close();
    }

    void print_values() const
    {
        for (int cell_ny = 0; cell_ny < grid_cell_ver_num; cell_ny++)
        {
            for (int cell_nx = 0; cell_nx < grid_cell_hor_num; cell_nx++)
            {
                for (int box = 0; box < box_num_per_cell; box++)
                {
                    f_vec_t params = get(cell_nx, cell_ny, box);
                    if (int(params[4]) != 0) // iou != 0
                    {
                        std::cout << "[ ";
                        for (int i = 0; i < PARAM_NUM_PER_BOX; i++)
                        {
                            std::cout << params[i] << " ";
                        }
                        std::cout << "], ";
                    }
                }
            }
            std::cout << std::endl;
        }
    }

    void normalize()
    {
        const int image_width = cell_width * grid_cell_hor_num;
        const int image_height = cell_height * grid_cell_ver_num;
        for (int cell_ny = 0; cell_ny < grid_cell_ver_num; cell_ny++)
        {
            for (int cell_nx = 0; cell_nx < grid_cell_hor_num; cell_nx++)
            {
                for (int box = 0; box < box_num_per_cell; box++)
                {
                    float_t *p = this->operator()(cell_nx, cell_ny, box);
                    p[0] /= 1;                   // cx
                    p[1] /= 1;                   // cy
                    p[2] /= cell_width;          // w
                    p[3] /= cell_height;         // h
                    p[4] *= (cell_width / 2.0f); // iou
                }
            }
        }
    }

    void denormalize()
    {
        const int image_width = cell_width * grid_cell_hor_num;
        const int image_height = cell_height * grid_cell_ver_num;
        for (int cell_ny = 0; cell_ny < grid_cell_ver_num; cell_ny++)
        {
            for (int cell_nx = 0; cell_nx < grid_cell_hor_num; cell_nx++)
            {
                for (int box = 0; box < box_num_per_cell; box++)
                {
                    float_t *p = this->operator()(cell_nx, cell_ny, box);
                    p[0] *= 1;                   // cx
                    p[1] *= 1;                   // cy
                    p[2] *= cell_width;          // w
                    p[3] *= cell_height;         // h
                    p[4] /= (cell_width / 2.0f); // iou
                }
            }
        }
    }

  private:
    float_t *operator()(const int nx, const int ny, const int nbox) const
    {
        assert(0 <= nx && nx < grid_cell_hor_num);
        assert(0 <= ny && ny < grid_cell_ver_num);
        assert(0 <= nbox && nbox < box_num_per_cell);

        return values + ny * row_size + nx * unit_size + nbox * PARAM_NUM_PER_BOX;
    }
};

#if YOLT_DEF_TEST

class TestYoltLabel
{

  public:
    TestYoltLabel()
    {
        ;
    }

    static void run()
    {
        YoltLabel test(416, 416, 26, 26, 2);
        show_vec(test.get(0, 0, 0));
        test.set(0, 0, 0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
        show_vec(test.get(0, 0, 0));

        show_vec(test.get(25, 25, 1));
        test.set(25, 25, 1, {1.0f, 2.0f, 2, 130, 132});
        show_vec(test.get(25, 25, 1));

        show_vec(test.get(13, 13, 0));
        test.set(13, 13, 0, {1.0f, 2.0f, 1.0f, 414.0f, 5.0f});
        show_vec(test.get(13, 13, 0));
    }

    static void show_vec(const std::vector<float_t> &vec)
    {
        for_each(vec.begin(), vec.end(),
                 [](const float_t v) -> void {
                     cout << v << " ";
                 });
        cout << endl;
    }
};

#endif

} // namespace nn_yolt
#endif // !YOLT_DEF_H